from langchain_openai.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, LLMRequestsChain
from langchain_community.tools import DuckDuckGoSearchRun


# Function to generate video script
def generate_script(prompt, video_length, creativity, api_key):
    # Template for generating 'Title'
    title_template = PromptTemplate(
        input_variables=['subject'],
        template='Please come up with a title for a YouTube video on the  {subject}.'
    )

    # Template for generating 'Video Script' using search engine
    script_template = PromptTemplate(
        input_variables=['title', 'DuckDuckGo_Search', 'duration'],
        template='Create a script for a YouTube video based on this title for me. TITLE: {title} of duration: {duration} minutes using this search data {DuckDuckGo_Search} '
    )

    # Setting up OpenAI LLM
    # llm = OpenAI(temperature=creativity, openai_api_key=api_key, model_name='gpt-3.5-turbo')
    llm = OpenAI(temperature=creativity, openai_api_key=api_key)

    # Creating chain for 'Title' & 'Video Script'
    title_chain = LLMChain(llm=llm, prompt=title_template, verbose=True)
    script_chain = LLMChain(llm=llm, prompt=script_template, verbose=True)

    # https://python.langchain.com/docs/modules/agents/tools/integrations/ddg
    # search = DuckDuckGoSearchRun()

    # Executing the chains we created for 'Title'
    title = title_chain.invoke({'subject': prompt}).get("text")

    # Executing the chains we created for 'Video Script' by taking help of search engine 'DuckDuckGo'
    # search_result = search.run(prompt)

    # DuckDuckGoSearchRun 有bug, 改用 bing + langchain request
    llm = OpenAI()
    template = """
    Extract the answer to the question '{query}' or say "not found" if the information is not available.
    {requests_result}
    """
    PROMPT = PromptTemplate(
        input_variables=["query", "requests_result"],
        template=template,
    )
    chain = LLMRequestsChain(llm_chain=LLMChain(llm=llm, prompt=PROMPT))

    inputs = {
        'query': prompt,
        'url': "https://cn.bing.com/search?q=" + prompt.replace(" ", "+"),
    }
    search_result = chain(inputs)
    # ------------------------------

    script = script_chain.run(title=title, DuckDuckGo_Search=search_result, duration=video_length)

    # Returning the output
    return search_result, title, script
